##################################################################
#
# PhyML-M3l landscape visualizer
#
# Victor Hanson-Smith
# victorhs@cs.uoregon.edu
# August 2009 
#
# DESCRIPTION:
# This script plots data from M3L trace lines that contain:
# . bfgs.c 1303
#
# This script parses the output from M3L, captures lines with performance data,
# and then plots the data as PDF files. If the PhyML execution was configured with
# low verbosity, then the output will not contain performance data.  In this case,
# this script will exit gracefully and not create PDF plots.
# 
# USAGE:
# python plot_performance.py <filepath>
# ... where <filepath> is the path to a text file capturing STDOUT from PhyML-M3l.  
#
#
#####################################################################

import os
import sys
import re
from dendropy import Tree
from argParser import *
ap = ArgParser(sys.argv)

try:
    filepath = sys.argv[1]
except IndexError:
    print "Yikes! I think you forgot to specify the filepath containing PhyML output."
    print "\nUsage: python plot_performance.py <filepath>"

try:
    fin = open(filepath, "r")
except IOError:
    print "Yikes! The file named ", filename, " does not exist."
    exit()

limit_plot = False
wantbranch = int( ap.getOptionalArg("--wantbranch") )
if wantbranch == False:
    wantbranch = -1
else:
    limit_plot = True

branches = []
plot1_xy = {} # key = branch id, value = hashtable of points

#
# Parse the PhyML output, grab relevant data.
#
# Looking for:
# . bfgs.c 1303 branch 368 l= -0.030000 lnl= -133178.045776
#
def get_all_points(fin):
    count_iterations = 0
    want_iter = 0
    last_hc_seconds = 0 # if we use a hyrbid approach (HC+STA), the clock gets reset after HC, so we need to adjust how we interpret the samples
    for line in fin.readlines():
        if line.__contains__("[Branch lengths"):
            count_iterations += 1
            return
        if line.__contains__(". bfgs.c 1303") and count_iterations == want_iter:
            tokens = line.split()
            this_branch = int(tokens[4])
            if False == branches.__contains__( this_branch ):
                branches.append( this_branch )
            if False == plot1_xy.__contains__( this_branch ):
                plot1_xy[this_branch] = {}
            plot1_xy[this_branch][float(tokens[6])] = float(tokens[8])
            #print this_branch, tokens[6], tokens[8]

def get_points_for_one_branch(fin):
    erg_length = None
    this_erg_length = None
    this_erg_lnl = None
    count_iterations = -1
    for line in fin.readlines():
        if line.__contains__(". debug bfgs.c 1303 - Printing the  likelihood landscape."):
            count_iterations += 1
        if line.__contains__("*l="):
            tokens = line.split()
            if erg_length == None:
                erg_length = tokens[11]
                erg_length = re.sub("\)", "", erg_length)
                erg_length = float( erg_length )
                this_erg_length = erg_length
            else:
                this_erg_length = tokens[11]
                this_erg_length = re.sub("\)", "", this_erg_length)
                this_erg_length = float( this_erg_length )
            this_erg_lnl = float(tokens[14])
            
        if line.__contains__(". bfgs.c 1303"):
            tokens = line.split()
            this_branch = int(tokens[4])
            if int(this_branch) == int(wantbranch) or int(wantbranch) == -1:
                # we plot multiple iterations of the same branch on one plot:
                this_branch = count_iterations
                if False == branches.__contains__( this_branch ):
                    branches.append( this_branch )
                if False == plot1_xy.__contains__( this_branch ):
                    plot1_xy[this_branch] = {}
                print this_erg_length
                length = (this_erg_length + float(tokens[6]))
                lnl_shift = this_erg_lnl + float(tokens[8])
                plot1_xy[this_branch][length] = lnl_shift
                print this_branch, length, lnl_shift


if limit_plot == False:
    get_all_points(fin)
else:
    get_points_for_one_branch(fin)


print "points = ", plot1_xy

colors = []
colors.append("black")
for i in range(1, (branches.__len__()+1)):
    colors.append(i.__str__())


#
# Sanity check
#
early_quit = False
if plot1_xy.__len__() < 1:
    print "I did not find any data for plot1(iterations vs. -lnL)"
    early_quit = True

def plot_in_r(points, output_filename_seed, title, xlab, ylab):
    maxx = None
    minx = None
    maxy = None
    miny = None

    branches = points.keys()
    branches.sort()

    # the ingredients for our CRAN script:
    plotstring = ""
    pointsstring = ""
    legendstring = "legend(\"bottomright\", c("
    for f in branches:
        legendstring += "\"" + f.__str__() + "\","
    legendstring = re.sub(",$", "", legendstring)
    legendstring += "), cex=0.8, col=c("
    for col in colors:
        legendstring += "\"" + col + "\","
    legendstring = re.sub(",$", "", legendstring)        
    legendstring += "), pch=15);\n"

    count = -1
    branchkeys = plot1_xy.keys()
    branchkeys.sort()
    for f in branchkeys:
        #print points[f]
        count += 1            
        x_sorted = points[f].keys()
        x_sorted.sort()
        print x_sorted
        
        string = "x" + count.__str__() + "<-c("
        for x in x_sorted:
            if maxx == None:
                maxx = x
            elif x > maxx:
                maxx = x
            if minx == None:
                minx = x
            elif x < minx:
                minx = x
            string += x.__str__() + ","   
        string = re.sub(",$", "", string)
        string += ");\n"
        
        string += "y" + count.__str__() + "<-c("
        for x in x_sorted:
            if miny == None:
                miny = points[f][x]
            elif points[f][x] < miny:
                miny = points[f][x]
            if maxy == None:
                maxy = points[f][x]
            elif points[f][x] > maxy:
                maxy = points[f][x]
            string += points[f][x].__str__() + ","
        string = re.sub(",$", "", string)
        string += ");\n"
        string += "points(x" + count.__str__() + ",y" + count.__str__() + ", col='" + colors[count] + "', type='l',lwd=2);\n"
        pointsstring += string
    
    plotstring = "x <-c(" + minx.__str__() + "," + maxx.__str__() + ");\n"
    plotstring += "y <-c(" + miny.__str__() + "," + maxy.__str__() + ");\n"
    plotstring += "plot(x, y, type='n', main='" + title + "', xlab='" + xlab + "', ylab='" + ylab + "');\n"
    
    fout_cran = open(output_filename_seed + ".cran", "w")
    fout_cran.write("pdf('" + output_filename_seed + ".pdf', width=4, height=7);\n")
    fout_cran.write(plotstring)
    fout_cran.write(pointsstring)
    #fout_cran.write(legendstring)
    fout_cran.write("dev.off();\n")
    fout_cran.close()
    
    os.system("r --no-save < " + output_filename_seed + ".cran")

plot_in_r(plot1_xy, "plot.landscape", "lnl landscape", "delta length", "lnL")
